from ultralytics import YOLO
import cv2
import io, os
from src.ms_server_api import MsServerApi

ms_api = MsServerApi()


class Detector:
    def __init__(self):
        # 模型路径在这里加载（只加载一次，避免重复初始化）
        self.fire_model = YOLO(r'weights/custom/fire.pt')
        self.helmet_model = YOLO(r'weights/custom/hat2.pt')
        self.person_model = YOLO(r'weights/yolo11n.pt')

    def predict_fire(self, source):
        """火焰识别"""
        predict_target = [0]  # head
        results = self.fire_model.predict(source=source, classes=predict_target, save=False)
        detected = self.check_have_target_class(predict_target, results)
        return results, detected

    def predict_person(self, source):
        """人类识别"""
        predict_target = [0]  # head
        results = self.person_model.predict(source=source, classes=predict_target, save=False)
        detected = self.check_have_target_class(predict_target, results)
        return results, detected

    def predict_helmet(self, source):
        """安全帽识别"""
        predict_target = [0]  # head
        results = self.helmet_model.predict(source=source, classes=predict_target, save=False)
        detected = self.check_have_target_class(predict_target, results)
        return results, detected

    def check_have_target_class(self, predict_target_classes: list[int], results) -> bool:
        """判断是否有目标类别"""
        detected = any(
            cls in predict_target_classes
            for r in results
            for cls in r.boxes.cls.int().tolist()  # 遍历这一帧中所有检测出的类别 ID
        )
        return detected

    def predict_hat2(self, source):
        """头盔识别 [0] head, 返回结果图片或视频"""
        results = self.helmet_model.predict(source=source,
                                            save=True,
                                            save_dir="runs/detect/predict",
                                            conf=0.5,
                                            classes=[0])

        # 判断是图片还是视频
        if source.lower().endswith(
                ('.mpo', '.tif', '.heic', '.webp', '.tiff', '.jpeg', '.dng', '.pfm', '.jpg', '.png', '.bmp')):
            result = results[0].plot()
            # numpy 转字节流
            _, buffer = cv2.imencode(".jpg", result)
            return ms_api.upload_bytes(buffer.tobytes(), "predict_res.jpg")
        elif source.lower().endswith(
                ('.m4v', '.wmv', '.webm', '.gif', '.mkv', '.asf', '.avi', '.mov', '.mp4', '.mpg', '.mpeg', '.ts')):
            # 上传预测结果视频
            save_dir = results[0].save_dir  # e.g. runs/detect/predict5
            video_name = os.path.basename(source)
            result_video_path = os.path.join(save_dir, video_name)

            with open(result_video_path, "rb") as f:
                video_bytes = f.read()
            return ms_api.upload_bytes(video_bytes, video_name)
        return None

    def predict(self, frame, target_classes):
        if target_classes is [0]:
            self.predict_helmet(frame)

    def predict_target(self, frame, predict_mark):
        """根据目标类型调用对应预测方法"""
        mapping = {
            'head': self.predict_helmet,
            'fire': self.predict_fire,
            'person': self.predict_person,
        }
        func = mapping.get(predict_mark)
        return func(frame) if func else (None, None)
